import os
import struct
import shutil
import numpy as np
from os import path
import sys
import utils

def release_label(filename,file):
	f=open(filename,'rb')
	index=0
	buf=f.read()
	f.close()
	magic,labels=struct.unpack_from('>II',buf,index)
	labelArr = [0] * labels
	index+=struct.calcsize('>II')
	for x in range(labels):
		labelArr[x]=int(struct.unpack_from('>B',buf,index)[0])
		index += struct.calcsize('>B')
	save=open(file,'w')
	save.write(','.join(map(lambda x: str(x), labelArr)))
	save.write('\n')
	save.close()
	return np.array(labelArr).reshape((1,labels))

def release_image(filename):
	f = open(filename,'rb')
	index = 0
	buf = f.read()
	f.close()
	magic,images,rows,columns=struct.unpack_from('>IIII',buf,index)
	index+=struct.calcsize('>IIII')
	images_counter=images
	format_str = '>%dB' % (rows*columns)
	images_bsize = struct.calcsize(format_str)
	samples=[]
	for i in range(images_counter):
		arr=struct.unpack_from(format_str,buf,index)
		samples.append(arr)
		index += images_bsize
	return np.array(samples)

def classify_trains_data(src,labels):
	cube=[[],[],[],[],[],[],[],[],[],[]]
	for i in range(labels.shape[1]):
		label=labels[0][i]
		image=src[i].reshape((1,src[i].size))
		cube[label].append(image[0])
	for j in range(len(cube)):
		print("the %d image counter is %d" % (j,len(cube[j])))
		utils.save_matrix(np.array(cube[j]),"models/original_train_matrix_"+str(j)+"_new")

def classify_test_data(src,labels):
	cube=[[],[],[],[],[],[],[],[],[],[]]
	for i in range(labels.shape[1]):
		label=labels[0][i]
		image=src[i].reshape((1,src[i].size))
		cube[label].append(image[0])
	for j in range(len(cube)):
		print("the %d image counter is %d" % (j,len(cube[j])))
		utils.save_matrix(np.array(cube[j]),"models/original_test_matrix_"+str(j)+"_new")

def run(src_file_name,lable_file_name,test_file_name,lable_test_name):
	lables=release_label(lable_file_name,"label_data/label_train.txt")
	imges=release_image(src_file_name)
	classify_trains_data(imges,lables)
	lables=release_label(lable_test_name,"label_data/label_test.txt")
	imges=release_image(test_file_name)
	classify_test_data(imges,lables)




if __name__ == "__main__":
	run("original_data/train-images.idx3-ubyte","original_data/train-labels.idx1-ubyte","original_data/t10k-images.idx3-ubyte","original_data/t10k-labels.idx1-ubyte")